import numpy as np
import taichi as ti
from JaccobiSolver import JaccobiSolver

# gui的分辨率
res = 256
# 时间步长
dt = 0.01
rho = 1
p_jacobi_iters = 4000  # 40 for a quicker but less accurate result
dx = 1 / res


ti.init(arch=ti.gpu)
print('Using jacobi iteration')

_velocities = ti.Vector.field(2, float, shape=(res, res))
_new_velocities = ti.Vector.field(2, float, shape=(res, res))
velocity_divs = ti.field(float, shape=(res, res))
dyes_divs = ti.field(float, shape=(res, res))
_pressures = ti.field(float, shape=(res, res))
_new_pressures = ti.field(float, shape=(res, res))
_dye_buffer = ti.Vector.field(3, float, shape=(res, res))
_new_dye_buffer = ti.Vector.field(3, float, shape=(res, res))

nParticles = 100000
particles = ti.Vector.field(2, dtype=ti.f32, shape=nParticles)
_new_dye_buffer = ti.Vector.field(3, float, shape=(res, res))

random_particles = ti.Vector.field(2, dtype=ti.f32, shape=nParticles//2)

# jaccobi_A = ti.field(dtype=ti.f32, shape=(res*res, res*res))
jaccobi_b = ti.field(dtype=ti.f32, shape=res*res)

K = ti.linalg.SparseMatrixBuilder(res*res, res*res, max_num_triplets=100)


@ti.kernel
def initialize():
    for i in particles:
        particles[i] = ti.random(), ti.random()
    for i in random_particles:
        random_particles[i] = ti.random(), ti.random()
    for i,j in _velocities:
        _velocities[i, j] = ti.random()*10,ti.random()*10
    # for i,j in _pressures:
    #     _pressures[i, j] = ti.random()


class TexPair:
    def __init__(self, cur, nxt):
        self.cur = cur
        self.nxt = nxt

    def swap(self):
        self.cur, self.nxt = self.nxt, self.cur


velocities_pair = TexPair(_velocities, _new_velocities)
pressures_pair = TexPair(_pressures, _new_pressures)
dyes_pair = TexPair(_dye_buffer, _new_dye_buffer)


@ti.func
def sample(qf, u, v):
    I = ti.Vector([int(u), int(v)])
    I = max(0, min(res - 1, I))
    return qf[I]


@ti.func
def lerp(vl, vr, frac):
    # frac: [0.0, 1.0]
    return vl + frac * (vr - vl)


@ti.func
def bilerp(vf, p):
    u, v = p
    s, t = u - 0.5, v - 0.5
    # floor
    iu, iv = ti.floor(s), ti.floor(t)
    # fract
    fu, fv = s - iu, t - iv
    a = sample(vf, iu, iv)
    b = sample(vf, iu + 1, iv)
    c = sample(vf, iu, iv + 1)
    d = sample(vf, iu + 1, iv + 1)
    return lerp(lerp(a, b, fu), lerp(c, d, fu), fv)


# 3rd order Runge-Kutta
@ti.func
def backtrace(vf: ti.template(), p, dt: ti.template()):
    v1 = bilerp(vf, p)
    p1 = p - 0.5 * dt * v1
    v2 = bilerp(vf, p1)
    p2 = p - 0.75 * dt * v2
    v3 = bilerp(vf, p2)
    p -= dt * ((2 / 9) * v1 + (1 / 3) * v2 + (4 / 9) * v3)
    return p


@ti.kernel
def advect(vf: ti.template(), qf: ti.template(), new_qf: ti.template()):
    for i, j in vf:
        p = ti.Vector([i, j]) + 0.5
        p = backtrace(vf, p, dt)
        new_qf[i, j] = bilerp(qf, p)


@ti.kernel
def Gravity(qf: ti.template(), new_qf: ti.template()):
    for i, j in qf:
        new_qf[i, j] = qf[i, j] + ti.Vector([0, -9.8*10])*dt


@ti.kernel
def divergence(vf: ti.template(), divs: ti.template()):
    for i, j in vf:
        vl = sample(vf, i - 1, j)
        vr = sample(vf, i + 1, j)
        vb = sample(vf, i, j - 1)
        vt = sample(vf, i, j + 1)
        vc = sample(vf, i, j)
        if i == 0:
            vl.x = -vc.x
        if i == res - 1:
            vr.x = -vc.x
        if j == 0:
            vb.y = -vc.y
        if j == res - 1:
            vt.y = -vc.y
        divs[i, j] = (vr.x - vl.x + vt.y - vb.y) * 0.5


@ti.kernel
def pressure_jacobi(pf: ti.template(), new_pf: ti.template(), divs: ti.template()):
    for i, j in pf:
        pl = sample(pf, i - 1, j)
        pr = sample(pf, i + 1, j)
        pb = sample(pf, i, j - 1)
        pt = sample(pf, i, j + 1)
        div = divs[i, j]
        new_pf[i, j] = (pl + pr + pb + pt - div) * 0.25


@ti.kernel
def subtract_gradient(vf: ti.template(), pf: ti.template()):
    for i, j in vf:
        pl = sample(pf, i - 1, j)
        pr = sample(pf, i + 1, j)
        pb = sample(pf, i, j - 1)
        pt = sample(pf, i, j + 1)
        vf[i, j] -= 0.5 * ti.Vector([pr - pl, pt - pb])


def solve_pressure_jacobi():
    for _ in range(p_jacobi_iters):
        pressure_jacobi(pressures_pair.cur, pressures_pair.nxt, velocity_divs)
        pressures_pair.swap()


def make_jaccobi_A(divs: ti.template()):
    for i, j in ti.ndrange(res, res):
        row = i * res + j
        center = 0.0
        if j != 0:
            K[row, row - 1] += 1.0
            center += 1.0
        if j != res - 1:
            K[row, row + 1] += 1.0
            center += 1.0
        if i != 0:
            K[row, row - res] += 1.0
            center += 1.0
        if i != res - 1:
            K[row, row + res] += 1.0
            center += 1.0
        K[row, row] += -center
    A = K.build()
    solver = ti.SparseSolver(solver_type="LLT")
    solver.analyze_pattern(A)
    solver.factorize(A)


@ti.kernel
def make_jaccobi_b(divs: ti.template()):
    for i,j in divs:
        row = i*res+j
        jaccobi_b[row] = divs[i,j]

def jacobi_solve(divs: ti.template(), res: ti.i32):
    make_jaccobi_b(velocity_divs)
    
    # js = JaccobiSolver(A=jaccobi_A, b=jaccobi_b, epochs=1000)
    # res = js.solve()
    # print(res)
    pass
    


@ti.kernel
def particle_move():
    for i in particles:
        x, y = particles[i] // dx
        grid_pos = ti.Vector([x, y])+0.5
        v = bilerp(velocities_pair.cur, grid_pos)
        # if particles[i][1]<1e-4 and v[1]<0:
        #     v = v[0]-0.8*v[1],-0.05*v[1]
        particles[i] += v * dt
    for i in random_particles:
        x, y = random_particles[i] // dx
        grid_pos = ti.Vector([x, y])+0.5
        v = bilerp(velocities_pair.cur, grid_pos)
        random_particles[i] += v * dt


def step():
    advect(velocities_pair.cur, velocities_pair.cur, velocities_pair.nxt)
    # advect(velocities_pair.cur, dyes_pair.cur, dyes_pair.nxt)
    velocities_pair.swap()
    # dyes_pair.swap()

    Gravity(velocities_pair.cur, velocities_pair.nxt)
    velocities_pair.swap()

    divergence(velocities_pair.cur, velocity_divs)

    solve_pressure_jacobi()
    jacobi_solve(velocity_divs,res)

    subtract_gradient(velocities_pair.cur, pressures_pair.cur)
    particle_move()


initialize()
make_jaccobi_A()
gui = ti.GUI('Stable Fluid', (res, res))
while gui.running:
    if gui.get_event(ti.GUI.PRESS):
        e = gui.event
        if e.key == ti.GUI.ESCAPE:
            break
    step()
    gui.circles(random_particles.to_numpy(), radius=1.5, color=0x00FF00)
    gui.circles(particles.to_numpy(), radius=1.5, color=0x0000FF)
    gui.show()
